import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import butter, cheby2, filtfilt, find_peaks
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score

# Folder containing all subject files
folder_path = 'TMCH_dataset/data/'  # Replace with your actual TMCH dataset folder path

# Function to normalize signal
def normalize_signal(signal):
    return (signal - np.min(signal)) / (np.max(signal) - np.min(signal))

# Function to apply Chebyshev Type II filter
def chebyshev_filter(signal, fs, order=4, cutoff=0.5, rs=40):
    nyquist = 0.5 * fs
    wn = cutoff / nyquist
    b, a = cheby2(order, rs, wn, btype='low', analog=False)
    return filtfilt(b, a, signal)

# IQR based outlier removal
def remove_outliers_iqr(signal):
    q1 = np.percentile(signal, 25)
    q3 = np.percentile(signal, 75)
    iqr = q3 - q1
    lower_bound = q1 - 1.5 * iqr
    upper_bound = q3 + 1.5 * iqr
    return signal[(signal >= lower_bound) & (signal <= upper_bound)]

# KSQI computation
def compute_ksqi(signal):
    mean = np.mean(signal)
    std = np.std(signal)
    if std == 0:
        return np.nan
    return np.mean(((signal - mean) / std) ** 4)

# Function to apply Butterworth low-pass filter
def butter_lowpass(cutoff, fs, order=4):
    nyquist = 0.5 * fs
    normal_cutoff = cutoff / nyquist
    b, a = butter(order, normal_cutoff, btype='low', analog=False)
    return b, a

def apply_butterworth_filter(signal, cutoff, fs, order=4):
    b, a = butter_lowpass(cutoff, fs, order=order)
    return filtfilt(b, a, signal)

# Parameters
sampling_rate = 100  # Hz for TMCH
cutoff_freq = 0.5  # Hz
window_durations = [10, 20, 30, 32, 64] + list(range(40, 130, 10))

# Output folder
output_folder = 'TMCH_Window_Size_Results'
os.makedirs(output_folder, exist_ok=True)

rmse_files = {wd: [] for wd in window_durations}
mae_files = {wd: [] for wd in window_durations}
r2_files = {wd: [] for wd in window_durations}

def save_results_to_csv(window_size, subject_results, metric_name):
    file_path = os.path.join(output_folder, f'Window_{window_size}s_{metric_name}.csv')
    df = pd.DataFrame(subject_results, columns=['FileName', metric_name])
    df.to_csv(file_path, index=False)
    print(f"Saved {metric_name} results for window size {window_size}s to {file_path}")

for file_name in os.listdir(folder_path):
    if file_name.endswith('.csv'):
        file_path = os.path.join(folder_path, file_name)
        data = pd.read_csv(file_path)

        file_identifier = os.path.splitext(file_name)[0]

        # Extract PPG_RED or PPG_IR and Actual RR
        ppg_signal = data['PPG_RED']  # or 'PPG_IR'
        actual_rr = data['Actual_RR']
        time = np.arange(len(ppg_signal)) / sampling_rate

        # Preprocessing
        ppg_signal = normalize_signal(ppg_signal)
        ppg_signal = chebyshev_filter(ppg_signal, sampling_rate)
        ppg_signal = remove_outliers_iqr(ppg_signal)
        ksqi_value = compute_ksqi(ppg_signal)
        ppg_signal = apply_butterworth_filter(ppg_signal, cutoff_freq, sampling_rate)

        for window_duration in window_durations:
            respiration_rates_calculated = []
            respiration_rates_actual = []

            for start_idx in range(0, len(ppg_signal), window_duration * sampling_rate):
                end_idx = start_idx + window_duration * sampling_rate
                if end_idx > len(ppg_signal):
                    break
                segment = ppg_signal[start_idx:end_idx]
                segment_rr = actual_rr[start_idx:end_idx]

                peaks, _ = find_peaks(segment, distance=sampling_rate)
                if len(peaks) > 1:
                    intervals = np.diff(peaks) / sampling_rate
                    est_rr = 60 / np.mean(intervals)
                else:
                    est_rr = np.nan
                respiration_rates_calculated.append(est_rr)
                respiration_rates_actual.append(np.mean(segment_rr))

            respiration_rates_calculated = np.array(respiration_rates_calculated)
            respiration_rates_actual = np.array(respiration_rates_actual)

            valid_idx = np.isfinite(respiration_rates_calculated) & np.isfinite(respiration_rates_actual)
            respiration_rates_calculated = respiration_rates_calculated[valid_idx]
            respiration_rates_actual = respiration_rates_actual[valid_idx]

            if len(respiration_rates_actual) > 1:
                rmse = np.sqrt(mean_squared_error(respiration_rates_actual, respiration_rates_calculated))
                mae = mean_absolute_error(respiration_rates_actual, respiration_rates_calculated)
                r2 = r2_score(respiration_rates_actual, respiration_rates_calculated) if np.var(respiration_rates_actual) > 0 else np.nan
            else:
                rmse = mae = r2 = np.nan

            rmse_files[window_duration].append((file_identifier, rmse))
            mae_files[window_duration].append((file_identifier, mae))
            r2_files[window_duration].append((file_identifier, r2))

for window_duration in window_durations:
    save_results_to_csv(window_duration, rmse_files[window_duration], 'RMSE')
    save_results_to_csv(window_duration, mae_files[window_duration], 'MAE')
    save_results_to_csv(window_duration, r2_files[window_duration], 'R2')

print("\nAll TMCH results saved.")